from PyQt5.QtWidgets import (QMainWindow, QVBoxLayout, QPushButton,
                             QFileDialog, QLabel, QWidget, QProgressBar, QMessageBox)
from PyQt5.QtCore import Qt
from docx_generator import DocxGenerator
from code_processor import CodeProcessor
import os

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.source_folder = set()
        self.setWindowTitle("软著代码整理工具")
        self.setFixedSize(400, 200)
        
        self.init_ui()
        self.code_processor = CodeProcessor()
        self.docx_gen = DocxGenerator()
        
    def init_ui(self):
        central = QWidget()
        layout = QVBoxLayout()
        
        self.label = QLabel("请选择源代码文件夹")
        self.label.setAlignment(Qt.AlignLeft)
        self.label.setWordWrap(True)
        
        self.btn_select = QPushButton("选择文件夹")
        self.btn_select.clicked.connect(self.select_folder)
        
        self.btn_generate = QPushButton("生成文档")
        self.btn_generate.clicked.connect(self.generate_doc)
        self.btn_generate.setEnabled(False)
        
        self.progress = QProgressBar()
        self.progress.setVisible(False)
        
        layout.addWidget(self.label)
        layout.addWidget(self.btn_select)
        layout.addWidget(self.btn_generate)
        layout.addWidget(self.progress)
        central.setLayout(layout)
        self.setCentralWidget(central)
    
    def select_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择源代码文件夹")
        if folder:
            self.source_folder.add(folder)
            text = "\n> ".join(self.source_folder)
            self.label.setText(f"已选择文件夹: \n> {text}")
            self.btn_generate.setEnabled(True)

    def generate_doc(self):
        if hasattr(self, 'source_folder'):
            self.output_path = QFileDialog.getSaveFileName(self, '保存Word文档', '', 'Word文件 (*.docx)')[0]
            if self.output_path:
                self.generate_word()
    def generate_word(self):
        self.progress.setVisible(True)
        self.progress.setValue(0)
        
        # 处理代码文件
        processed_code = self.code_processor.process_folder(self.source_folder)
        # 生成Word文档
        self.docx_gen.generate(processed_code, self.output_path, self.update_progress)
        
        self.progress.setValue(100)
        self.label.setText("文档生成完成！")
        self.source_folder.clear()
        QMessageBox.information(self, "提示", "文档生成完成！")

    def update_progress(self, progress):
        self.progress.setValue(progress)